package com.example.mybatis_demo.datasource;

import com.example.mybatis_demo.po.User;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.util.Properties;

@Intercepts({
        @Signature(type = Executor.class, method = "query",
                args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class MybatisDatasourceInterceptor implements Interceptor {
    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        // 方法的参数
        Object[] args = invocation.getArgs();
        Object parameterObject = args[1];
        if (parameterObject instanceof User) {
            User user = (User) parameterObject;
            final int shardFlag = user.getShardFlag();
            final DataSourceConstants value = DataSourceConstants.values()[shardFlag & 1];
            DynamicDataSourceContextHolder.setContextKey(value);
        }
        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.removeContextKey();
        }
    }

    @Override
    public Object plugin(Object target) {
        return Interceptor.super.plugin(target);
    }

    @Override
    public void setProperties(Properties properties) {
        Interceptor.super.setProperties(properties);
    }
}
